package com.lhb.aicodemother.ai.core.handler;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.lhb.aicodemother.ai.core.builder.VueProjectBuilder;
import com.lhb.aicodemother.ai.model.enums.StreamMessageTypeEnum;
import com.lhb.aicodemother.ai.model.message.AiResponseMessage;
import com.lhb.aicodemother.ai.model.message.StreamMessage;
import com.lhb.aicodemother.ai.model.message.ToolExecutionMessage;
import com.lhb.aicodemother.ai.model.message.ToolRequestMessage;
import com.lhb.aicodemother.constant.AppConstant;
import com.lhb.aicodemother.model.entity.User;
import com.lhb.aicodemother.model.enums.ChatHistoryMessageTypeEnum;
import com.lhb.aicodemother.service.ChatHistoryService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.HashSet;
import java.util.Set;

/*
* JSON消息流处理器
* */
@Slf4j
@Component
public class JsonMessageStreamHandler {

    @Resource
    private VueProjectBuilder vueProjectBuilder;

    /*
    * 处理tokenStream
    * 解析JSON消息并重组为完整的响应格式
    * */
    public Flux<String> handle(Flux<String> orginFlux,
                               ChatHistoryService chatHistoryService,
                               long appId, User loginUser){
        // 收集数据用于生成后端记忆格式
        StringBuilder chatHistoryStringBuilder = new StringBuilder();
        // 用于跟踪已经见过的工具Id，判断是否为第一次使用
        Set<String> seenToolIds = new HashSet<String>();

        return orginFlux
                .map(chunk -> {
                    // 解析每个JSON消息块
                    return handleJsonMessageChunk(chunk,chatHistoryStringBuilder,seenToolIds);
                })
                .filter(StrUtil::isNotEmpty)
                .doOnComplete(() -> {
                    // 流式响应完成后，添加AI消息到对话历史
                    String aiResponse = chatHistoryStringBuilder.toString();
                    chatHistoryService.addChatMessage(appId,aiResponse, ChatHistoryMessageTypeEnum.AI.getValue(),  loginUser.getId());

                    // 异步构建Vue项目
                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + "/vue_project_" + appId;
                    vueProjectBuilder.buildProjectAsync(projectPath);
                })
                .doOnError(error -> {
                    // 如果AI回复失败，也要记录错误消息
                    String errorMessage = "AI回复失败：" + error.getMessage();
                    chatHistoryService.addChatMessage(appId,errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                });

    }

    /*
    * 解析并收集TokenStream数据
    * */
    private String handleJsonMessageChunk(String chunk, StringBuilder chatHistoryStringBuilder, Set<String> seenToolIds){
        // 解析JSON
        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (typeEnum) {
            case AI_RESPONSE -> {
                AiResponseMessage aiResponseMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                String data = aiResponseMessage.getData();
                // 直接拼接响应
                chatHistoryStringBuilder.append(data);
                return data;
            }
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
                String toolId = toolRequestMessage.getId();
                // 检查是否第一次看到这个工具ID
                if(toolId != null && !seenToolIds.contains(toolId)){
                    // 第一次调用这个工具，记录ID并返回工具信息
                    seenToolIds.add(toolId);
                    return "\n\n[选择工具] 写入文件\n\n";
                }else {
                    // 不是第一次调用这个工具，直接返回空
                    return "";
                }

            }
            case TOOL_EXECUTED -> {
                ToolExecutionMessage  toolExecutionMessage = JSONUtil.toBean(chunk, ToolExecutionMessage.class);
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutionMessage.getArguments());
                String relativeFilePath = jsonObject.getStr("relativeFilePath");
                String suffix = FileUtil.getSuffix(relativeFilePath);
                String content = jsonObject.getStr("content");
                String result = String.format("""
                        [工具调用] 写入文件 %s
                        ```%s
                        %s
                        """,relativeFilePath,suffix,content);
                // 输出前端和要持久化的内容
                String output = String.format("\n\n%s\n\n",result);
                chatHistoryStringBuilder.append(output);
                return output;
            }
            default -> {
                log.error("不支持的消息类型：{}",typeEnum);
                return "";
            }
        }
    }


}
